import os

import torch
from diffusers import StableDiffusion3Pipeline, AutoencoderKL, UNet2DConditionModel, DiffusionPipeline, \
    StableDiffusionXLPipeline
from transformers import CLIPTextModel

os.environ["http_proxy"] = "http://192.168.3.116:7890/"
os.environ["https_proxy"] = "http://192.168.3.116:7890/"

model_key = "stabilityai/stable-diffusion-xl-base-1.0"
# AutoencoderKL main
# CLIPTextModel main
# revision = "fp16"
revision = None

# stable_diffusion_pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16)
# pipe.to("cuda")

# stable_diffusion_pipe.scheduler.compatibles

# AutoencoderKL.from_pretrained(model_key, subfolder="vae", revision=revision,
#                                                  torch_dtype=torch.float16)
# CLIPTextModel.from_pretrained(model_key, subfolder="text_encoder", revision=revision,
#                                                           torch_dtype=torch.float16)

# DiffusionPipeline.from_pretrained(model_key, variant=revision,
#                                                          torch_dtype=torch.float16)

StableDiffusionXLPipeline.from_pretrained(model_key, variant=revision,
                                                         torch_dtype=torch.float16)

# print(stable_diffusion_pipe.scheduler.compatibles)